from models.resnet50 import *
from core.range_deduction import *

# 创建RangeDeduction实例
rd = RangeDeduction()

# 添加卷积层测试
conv_layer = LayerInfo(
    layer_id='conv1',
    layer_type='conv2d',
    kernel_size=(3,3),
    stride=(2,2),
    padding=(1,1),
    dilation=(2,2),
    input_shape=(0,1079,0,1919),
    output_shape=(0,538,0,958)
)
rd.register_layer(conv_layer)

# 添加池化层测试
pool_layer = LayerInfo(
    layer_id='pool1',
    layer_type='pool2d',
    kernel_size=(2,2),
    stride=(2,2),
    padding=(0,0),
    dilation=(1,1),
    input_shape=(0,538,0,958),
    output_shape=(0,268,0,478)
)
rd.register_layer(pool_layer)

# # 执行范围计算
# print('\n=== 卷积层测试 ===')
# conv_range = rd.calculate_input_range('conv1', (0, 538, 0, 958))
# print('最终输入范围:', conv_range)

# print('\n=== 池化层测试 ===')
# pool_range = rd.calculate_input_range('pool1', (0, 268, 0, 478))
# print('最终输入范围:', pool_range)

rd.add_layer_dependency('conv1', 'pool1')

# dependent_ranges = rd.get_dependent_ranges('pool1', (0, 268, 0, 478))
# print("\n单路径测试结果：")
# for layer_id, input_range in dependent_ranges.items():
#     print(f"Layer {layer_id} requires output range {input_range}")

minimal_ranges, intermediate_ranges = rd.AIR('conv1','pool1',(0, 268, 0, 478))
print("\nAIR测试结果：")
print("Minimal Ranges:", minimal_ranges)
print("Intermediate Ranges:", intermediate_ranges)

# 提取并注册层信息
def extract_and_register_layers(model, rd):
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d):
            rd.register_layer(LayerInfo(
                layer_id=name,
                layer_type='conv2d',
                kernel_size=module.kernel_size,
                stride=module.stride,
                padding=module.padding,
                dilation=module.dilation,
                input_shape=(3, 1080, 1920),  # 假设输入图像为1080P
                output_shape=(module.out_channels, 1080 // module.stride[0], 1920 // module.stride[1])
            ))
        elif isinstance(module, nn.MaxPool2d):
            rd.register_layer(LayerInfo(
                layer_id=name,
                layer_type='pool2d',
                kernel_size=(module.kernel_size, module.kernel_size),
                stride=(module.stride, module.stride),
                padding=(module.padding, module.padding),
                dilation=(1, 1),
                input_shape=(3, 1080, 1920),  # 假设输入图像为1080P
                output_shape=(3, 1080 // module.stride, 1920 // module.stride)
            ))
        elif isinstance(module, nn.ReLU):
            rd.register_layer(LayerInfo(
                layer_id=name,
                layer_type='relu',
                kernel_size=(1, 1),
                stride=(1, 1),
                padding=(0, 0),
                dilation=(1, 1),
                input_shape=(3, 1080, 1920),  # 假设输入图像为1080P
                output_shape=(3, 1080, 1920)
            ))
        elif isinstance(module, nn.BatchNorm2d):
            rd.register_layer(LayerInfo(
                layer_id=name,
                layer_type='bn2d',
                kernel_size=(1, 1),
                stride=(1, 1),
                padding=(0, 0),
                dilation=(1, 1),
                input_shape=(3, 1080, 1920),  # 假设输入图像为1080P
                output_shape=(3, 1080, 1920)
            ))

def add_layer_dependencies(rd):
    # 添加初始层的依赖关系
    rd.add_layer_dependency('conv1', 'bn1')
    rd.add_layer_dependency('bn1', 'relu')
    rd.add_layer_dependency('relu', 'maxpool')
    
    # 添加layer1的依赖关系
    rd.add_layer_dependency('maxpool', 'layer1')
    for i in range(3):
        rd.add_layer_dependency(f'layer1.{i}', f'layer1.{i+1}' if i < 2 else 'layer2')
    
    # 添加layer2的依赖关系
    rd.add_layer_dependency('layer1', 'layer2')
    for i in range(4):
        rd.add_layer_dependency(f'layer2.{i}', f'layer2.{i+1}' if i < 3 else 'layer3')
    
    # 添加layer3的依赖关系
    rd.add_layer_dependency('layer2', 'layer3')
    for i in range(6):
        rd.add_layer_dependency(f'layer3.{i}', f'layer3.{i+1}' if i < 5 else 'layer4')
    
    # 添加layer4的依赖关系
    rd.add_layer_dependency('layer3', 'layer4')
    for i in range(3):
        rd.add_layer_dependency(f'layer4.{i}', f'layer4.{i+1}' if i < 2 else 'avgpool')
    
    # 添加全局平均池化和分类器的依赖关系
    rd.add_layer_dependency('layer4', 'avgpool')
    rd.add_layer_dependency('avgpool', 'fc')

def test_single_path():
    # 初始化ResNet50模型
    model = ResNet50(num_classes=1000)
    # 创建RangeDeduction实例
    single_rd = RangeDeduction()
    # 提取并注册层信息
    extract_and_register_layers(model, single_rd)
    # 添加依赖关系
    add_layer_dependencies(single_rd)
    # 测试单一路径
    output_range = (0, 360, 0, 640)
    dependent_ranges = single_rd.get_dependent_ranges('maxpool', output_range)
    print("\n单一路径测试结果：")
    for layer_id, input_range in dependent_ranges.items():
        print(f"Layer {layer_id} requires input range {input_range}")

# 新增多祖先路径合并验证（论文步骤5）
def test_multi_ancestor_merge():
    # 创建新测试实例
    test_rd = RangeDeduction()

    # 注册多分支层结构
    layers = [
        ('conv1', 'conv2d'), ('branch1', 'conv2d'), ('branch2', 'conv2d'),
        ('merge', 'conv2d'), ('output', 'conv2d')
    ]
    for lid, ltype in layers:
        test_rd.register_layer(LayerInfo(
            layer_id=lid,
            layer_type=ltype,
            kernel_size=(3,3),
            stride=(1,1),
            padding=(1,1),
            dilation=(1,1),
            input_shape=(3, 512, 512),
            output_shape=(64, 512, 512)
        ))
    
    # 设置依赖关系
    test_rd.add_layer_dependency('conv1', 'branch1')
    test_rd.add_layer_dependency('conv1', 'branch2')
    test_rd.add_layer_dependency('branch1', 'merge')
    test_rd.add_layer_dependency('branch2', 'merge')
    test_rd.add_layer_dependency('merge', 'output')
    
    # 计算合并范围
    ranges = test_rd.get_dependent_ranges('output', (100, 200, 100, 200))
    
    # 调试输出
    print('\n调试信息：')
    for layer in ['output', 'merge', 'branch1', 'branch2']:
        print(f'{layer}层输入范围: {ranges.get(layer, "未找到")}')
    
    # 验证分支合并结果
    assert ranges['merge'] == (100, 200, 100, 200), "合并层范围错误"
    assert ranges['branch1'] == ranges['branch2'], "祖先路径范围应相同"
    print("\n多祖先路径合并验证通过")

# # 执行单一路径测试
# test_single_path()

# # 执行多分支测试
# test_multi_ancestor_merge()